#include <stdio.h>
#include <stdlib.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <time.h>
#include "prof.h"

/* CUDA错误处理宏 */
#define CHECK_CUDA_ERROR(call) { \
    cudaError_t err = call; \
    if (err != cudaSuccess) { \
        fprintf(stderr, "CUDA错误: %s\n", cudaGetErrorString(err)); \
        exit(EXIT_FAILURE); \
    } \
}

/* cuBLAS错误处理宏 */
#define CHECK_CUBLAS_ERROR(call) { \
    cublasStatus_t err = call; \
    if (err != CUBLAS_STATUS_SUCCESS) { \
        fprintf(stderr, "cuBLAS错误: %d\n", err); \
        exit(EXIT_FAILURE); \
    } \
}

/* 卡尔曼滤波内核函数（标量版本） */
__global__ void kalman_filter_kernel(
    const float *z,      // 观测数据，[m, T]
    const float F,       // 状态转移系数
    const float H,       // 观测系数
    const float Q,       // 过程噪声方差
    const float R,       // 观测噪声方差
    float *x_hat,        // 状态估计，[m, T]
    float *P,            // 协方差，[m, T]
    int T,               // 时间步数
    int m                // 组数
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x; // 线程索引
    if (idx < m) {
        float x = 0.0f; // 初始状态
        float P_val = 1.0f; // 初始协方差
        for (int t = 0; t < T; t++) {
            // 预测步骤
            float x_pred = F * x; // 状态预测
            float P_pred = F * P_val * F + Q; // 协方差预测
            // 更新步骤
            float K = P_pred * H / (H * P_pred * H + R); // 卡尔曼增益
            x = x_pred + K * (z[idx * T + t] - H * x_pred); // 状态更新
            P_val = (1 - K * H) * P_pred; // 协方差更新
            // 保存结果
            x_hat[idx * T + t] = x;
            P[idx * T + t] = P_val;
        }
    }
}

int main(int argc, char** argv) {
    int n, m, T;

    // 获取运行次数n
    if (argc != 2) {
        fprintf(stderr, "用法: %s <运行次数>\n", argv[0]);
        return 1;
    }
    n = atoi(argv[1]);
    if (n <= 0) {
        fprintf(stderr, "运行次数必须为正整数。\n");
        return 1;
    }

    profInit();//初始化
    
    profSetSectName(1, "kalman htod 时间");
    profSetSectName(2, "kalman compute_time");
    profSetSectName(3, "kalman dtoh 时间");

    // 信号组数
    m = 10000;
    // 时间步数 (T)
    T = 4096;

    /* 定义系统模型参数（标量） */
    float F = 1.0f; // 状态转移系数
    float H = 1.0f; // 观测系数
    float Q = 0.1f; // 过程噪声方差
    float R = 1.0f; // 观测噪声方差

    /* 分配主机内存 */
    float *h_z = (float*)malloc(m * T * sizeof(float)); // 观测数据
    float *h_x_hat = (float*)malloc(m * T * sizeof(float)); // 状态估计
    float *h_P = (float*)malloc(m * T * sizeof(float)); // 协方差

    /* 生成随机观测数据 */
    srand(time(NULL));
    for (int i = 0; i < m * T; i++) {
        h_z[i] = (float)rand() / RAND_MAX * 10.0f; // 随机观测值，范围[0,10]
    }

    /* 分配设备内存 */
    float *d_z, *d_x_hat, *d_P;
    CHECK_CUDA_ERROR(cudaMalloc(&d_z, m * T * sizeof(float)));
    CHECK_CUDA_ERROR(cudaMalloc(&d_x_hat, m * T * sizeof(float)));
    CHECK_CUDA_ERROR(cudaMalloc(&d_P, m * T * sizeof(float)));

    /* 运行n次 */
    printf("\n运行卡尔曼滤波...\n");

    profEnterSect(0);
    for (int i = 0; i < n; i++) {
        /* 记录htod时间 */
        profEnterSect(1);
        CHECK_CUDA_ERROR(cudaMemcpy(d_z, h_z, m * T * sizeof(float), cudaMemcpyHostToDevice));
        profLeaveSect(1);

        /* 记录计算时间 */
        profEnterSect(2);
        int threads_per_block = 256; // 每个线程块256个线程
        int blocks = (m + threads_per_block - 1) / threads_per_block; // 计算所需块数
        kalman_filter_kernel<<<blocks, threads_per_block>>>(d_z, F, H, Q, R, d_x_hat, d_P, T, m);
        CHECK_CUDA_ERROR(cudaDeviceSynchronize());
        profLeaveSect(2);

        /* 记录dtoh时间 */
        profEnterSect(3);
        CHECK_CUDA_ERROR(cudaMemcpy(h_x_hat, d_x_hat, m * T * sizeof(float), cudaMemcpyDeviceToHost));
        CHECK_CUDA_ERROR(cudaMemcpy(h_P, d_P, m * T * sizeof(float), cudaMemcpyDeviceToHost));
        profLeaveSect(3);

    }
    profLeaveSect(0);
    profPrintInfo();

    

    /* 清理资源 */
    free(h_z);
    free(h_x_hat);
    free(h_P);
    CHECK_CUDA_ERROR(cudaFree(d_z));
    CHECK_CUDA_ERROR(cudaFree(d_x_hat));
    CHECK_CUDA_ERROR(cudaFree(d_P));

    return 0;
}
